{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netCDF4 import Dataset\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.cm import get_cmap\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "import matplotlib.gridspec as gridspec\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords, destagger\n",
    "import glob\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import cmocean\n",
    "import pyresample\n",
    "from scipy.spatial import cKDTree\n",
    "from scipy.signal import argrelextrema\n",
    "import copy\n",
    "import Nio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = \"mynn\"\n",
    "dom1 = \"d01\"\n",
    "dom2 = \"d02\"\n",
    "active_dom = \"d02\"\n",
    "sens = \"no_fractional_seaice\"\n",
    "day = \"12\"\n",
    "#root = \"/glade/scratch/tjuliano/doe_comble/new_lam_runs/WRF/WRF-ASR-CAO/test/\"\n",
    "root = \"/glade/campaign/uwyo/wyom0122/lam/new_runs/\"\n",
    "cases = [\"031220_no_edmf\",\"031220_mynn_l3\",\"031220\",\"031220_ysu\"]\n",
    "rst = \"/RESTART_D\"\n",
    "if active_dom == \"d01\":\n",
    "    plt_titles = [r\"$\\Delta$x=3km, level 2.5\", \"$\\Delta$x=3km, level 3\", \"$\\Delta$x=3km, level 2.5 EDMF\", \"$\\Delta$x=3km, YSU\"]\n",
    "    plt_titles2 = [r\"$\\Delta$x=3km\"\"\\n\"\"Level 2.5\", \"$\\Delta$x=3km\"\"\\n\"\"Level 3\", \"$\\Delta$x=3km\"\"\\n\"\"Level 2.5 EDMF\", \"$\\Delta$x=3km\"\"\\n\"\"YSU\"]\n",
    "    dx = 3000.\n",
    "else:\n",
    "    plt_titles = [r\"$\\Delta$x=1km, level 2.5\", \"$\\Delta$x=1km, level 3\", \"$\\Delta$x=1km, level 2.5 EDMF\", \"$\\Delta$x=1km, YSU\"]\n",
    "    plt_titles2 = [r\"$\\Delta$x=1km\"\"\\n\"\"Level 2.5\", \"$\\Delta$x=1km\"\"\\n\"\"Level 3\", \"$\\Delta$x=1km\"\"\\n\"\"Level 2.5 EDMF\", \"$\\Delta$x=1km\"\"\\n\"\"YSU\"]\n",
    "    dx = 1000.\n",
    "savedir = \"figs_new_nz136/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "lwprng = np.arange(-30,80,5)\n",
    "times = [9,11,13,19,25,31,37,43,49]\n",
    "time_str = ['28_1030z','28_1130z','28_1230z','28_1530z','28_1830z','28_2130z','29_0030z','29_0330z','29_0630z']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "lev = 10\n",
    "skipx = 6\n",
    "skipy = 6\n",
    "skipx2 = skipx*3\n",
    "skipy2 = skipy*3\n",
    "scale = 50\n",
    "barbw = 0.04\n",
    "g = 9.81"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Nio.open_file(root + cases[0] + \"/wrfinput_\" + dom1,format='netcdf')\n",
    "cosalpha = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha = s1.variables['SINALPHA'][-1,:,:]\n",
    "s1.close()\n",
    "\n",
    "s1 = Nio.open_file(root + cases[3] + \"/wrfinput_\" + dom2,format='netcdf')\n",
    "cosalpha2 = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha2 = s1.variables['SINALPHA'][-1,:,:]\n",
    "s1.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = 'cells'\n",
    "if region == 'cells':\n",
    "###### OPEN CELLS\n",
    "#\tlon_arr = [12.0, 12.0, 16.0, 16.0]\n",
    "#\tlat_arr = [69.5, 71.5, 71.5, 69.5]\n",
    "\tlon_arr = [12.5, 12.5, 16.5, 16.5]\n",
    "\tlat_arr = [69.0, 71.0, 71.0, 69.0]\n",
    "elif region == 'rolls':\n",
    "###### ROLLS/CLOSED CELLS\n",
    "##        lon_arr = [7.5,  7.5, 12.5, 12.5]\n",
    "##        lat_arr = [71.5, 73.0, 73.0, 71.5]\n",
    "\tlon_arr = [8.0, 8.0, 12.0, 12.0]\n",
    "\tlat_arr = [74.5, 76.5, 76.5, 74.5]\n",
    "else:\n",
    "        print ('You are an idiot!')\n",
    "        \n",
    "#lat2 = getvar(s1,\"lat\",meta=False)\n",
    "#lon2 = getvar(s1,\"lon\",meta=False)\n",
    "#hgt = getvar(s1,\"HGT\",meta=False)\n",
    "#lat_d01 = lat2\n",
    "#lon_d01 = lon2\n",
    "#hgt_d01 = hgt\n",
    "ll_lon = lon_arr[0]\n",
    "ll_lat = lat_arr[0]\n",
    "ul_lon = lon_arr[0]\n",
    "ul_lat = lat_arr[1]\n",
    "ur_lon = lon_arr[2]\n",
    "ur_lat = lat_arr[1]\n",
    "lr_lon = lon_arr[2]\n",
    "lr_lat = lat_arr[0]\n",
    "# center point of overall domain\n",
    "ref_lat = (ll_lat + ur_lat)/2.\n",
    "#ref_lon = -(abs(ll_lon) + abs(ur_lon))/2.\n",
    "ref_lon = (ll_lon + ur_lon)/2.\n",
    "######################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files_wind0 = sorted(glob.glob(root + cases[0] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind1 = sorted(glob.glob(root + cases[1] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind2 = sorted(glob.glob(root + cases[2] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind3 = sorted(glob.glob(root + cases[3] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "\n",
    "files_cloud0 = sorted(glob.glob(root + cases[0] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "files_cloud1 = sorted(glob.glob(root + cases[1] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "files_cloud2 = sorted(glob.glob(root + cases[2] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "files_cloud3 = sorted(glob.glob(root + cases[3] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "\n",
    "times = [12]\n",
    "first_pass1 = 0\n",
    "do_qt_flux = False\n",
    "for i in times: #np.arange(12,len(files_cloud0),6):\n",
    "\tfor j in np.arange(len(cases)):\n",
    "\t\tprint ('Reading file...')\n",
    "\t\tif j == 0:\n",
    "\t\t\ts1 = Nio.open_file(files_wind0[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud0[i],format='netcdf')\n",
    "\t\t\tprint (files_wind0[i])\n",
    "\t\telif j == 1:\n",
    "\t\t\ts1 = Nio.open_file(files_wind1[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud1[i],format='netcdf')\n",
    "\t\t\tprint (files_wind1[i])\n",
    "\t\telif j == 2:\n",
    "\t\t\ts1 = Nio.open_file(files_wind2[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud2[i],format='netcdf')\n",
    "\t\t\tprint (files_wind2[i])\n",
    "\t\telif j == 3:\n",
    "\t\t\ts1 = Nio.open_file(files_wind3[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud3[i],format='netcdf')\n",
    "\t\t\tprint (files_wind3[i])\n",
    "\t\telif j == 4:\n",
    "\t\t\ts1 = Nio.open_file(files_wind4[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud4[i],format='netcdf')\n",
    "\t\t\tprint (files_wind4[i])\n",
    "\t\tif j == 0:\n",
    "\t\t\ts3 = Nio.open_file(root + cases[0] + \"/wrfinput_\" + active_dom,format='netcdf')\n",
    "\t\t\tlat = getvar(s3,\"lat\",meta=False)\n",
    "\t\t\tlon = getvar(s3,\"lon\",meta=False)\n",
    "\t\tseaice = getvar(s3,\"SEAICE\",meta=False)\n",
    "\t\tu = getvar(s1,\"ua\",units=\"m s-1\",meta=False)\n",
    "\t\tv = getvar(s1,\"va\",units=\"m s-1\",meta=False)\n",
    "\t\tw = getvar(s1,\"wa\",units=\"m s-1\",meta=False)\n",
    "\t\tt2 = getvar(s1,\"T2\",meta=False)\n",
    "\t\tz = getvar(s1,\"z\", units=\"m\",meta=False)\n",
    "\t\tpblh = getvar(s1,\"PBLH\",meta=False)\n",
    "\t\tqv = getvar(s2,\"QVAPOR\",meta=False)\n",
    "\t\t#qc = getvar(s2,\"QCLOUD\",meta=False)\n",
    "\t\t#qi = getvar(s2,\"QICE\",meta=False)\n",
    "\t\t#qs = getvar(s2,\"QSNOW\",meta=False)\n",
    "\t\t#qg = getvar(s2,\"QGRAUP\",meta=False)\n",
    "        \n",
    "\t\tif do_qt_flux == True:\n",
    "\t\t\tqt = qv#+qc+qi+qs+qg\n",
    "        \n",
    "\t\t\tp = getvar(s1,\"P\",meta=False)\n",
    "\t\t\tpb = getvar(s1,\"PB\",meta=False)\n",
    "        \n",
    "\t\t\tptot = (p+pb)/100.\n",
    "\n",
    "\t\t\tprint ('Interpolating...')\n",
    "\t\t\tp1 = 975.\n",
    "\t\t\tp2 = 925.\n",
    "\t\t\tdp = (p1-p2)*100.\n",
    "\t\t\tui1 = interplevel(u, ptot, p1, meta=False)\n",
    "\t\t\tui2 = interplevel(u, ptot, p2, meta=False)\n",
    "\t\t\tvi1 = interplevel(v, ptot, p1, meta=False)\n",
    "\t\t\tvi2 = interplevel(v, ptot, p2, meta=False)\n",
    "\t\t\tqt1 = interplevel(qt, ptot, p1, meta=False)\n",
    "\t\t\tqt2 = interplevel(qt, ptot, p2, meta=False)\n",
    "\t\t\twi = interplevel(w, ptot, (p1+p2)/2.)\n",
    "        \n",
    "\t\t\tif active_dom == \"d01\":\n",
    "\t\t\t\tuearth1 = ui1*cosalpha - vi1*sinalpha\n",
    "\t\t\t\tvearth1 = vi1*cosalpha + ui1*sinalpha\n",
    "\t\t\t\tuearth2 = ui2*cosalpha - vi2*sinalpha\n",
    "\t\t\t\tvearth2 = vi2*cosalpha + ui2*sinalpha\n",
    "\t\t\telse:\n",
    "\t\t\t\tuearth1 = ui1*cosalpha2 - vi1*sinalpha2\n",
    "\t\t\t\tvearth1 = vi1*cosalpha2 + ui1*sinalpha2\n",
    "\t\t\t\tuearth2 = ui2*cosalpha2 - vi2*sinalpha2\n",
    "\t\t\t\tvearth2 = vi2*cosalpha2 + ui2*sinalpha2\n",
    "\n",
    "\t\t#uearth1b = destagger(uearth1,0)\n",
    "\t\t#vearth1b = destagger(vearth1,1)\n",
    "\t\t#uearth2b = destagger(uearth2,0)\n",
    "\t\t#vearth2b = destagger(vearth2,1)\n",
    "        \n",
    "\t\t#print (np.shape(qv1),np.shape(qv2),np.shape(uearth1b),np.shape(uearth2b),np.shape(vearth1b),np.shape(vearth2b))\n",
    "        \n",
    "\t\t\tconv1 = ((qt1[:,1:]*uearth1[:,1:])-(qt1[:,0:-1]*uearth1[:,0:-1]))/dx + ((qt1[:,1:]*vearth1[:,1:])-(qt1[:,0:-1]*vearth1[:,0:-1]))/dx\n",
    "\t\t\tconv2 = ((qt2[:,1:]*uearth2[:,1:])-(qt2[:,0:-1]*uearth2[:,0:-1]))/dx + ((qt2[:,1:]*vearth2[:,1:])-(qt2[:,0:-1]*vearth2[:,0:-1]))/dx\n",
    "\t\t\tvimfc = -(1./g)*(conv2-conv1)/dp\n",
    "            \n",
    "\t\telse:\n",
    "\t\t\tprint ('Interpolating...')\n",
    "\t\t\thgt = pblh/2.\n",
    "\t\t\thgt2 = pblh/10.\n",
    "\t\t\thgt2 = 100.\n",
    "\t\t\tui = interplevel(u, z, hgt2, meta=False)\n",
    "\t\t\tvi = interplevel(v, z, hgt2, meta=False)\n",
    "\t\t\twi = interplevel(w, z, hgt)\n",
    "\n",
    "\t\t\tif active_dom == \"d01\":\n",
    "\t\t\t\tuearth = ui*cosalpha - vi*sinalpha\n",
    "\t\t\t\tvearth = vi*cosalpha + ui*sinalpha\n",
    "\t\t\telse:\n",
    "\t\t\t\tuearth = ui*cosalpha2 - vi*sinalpha2\n",
    "\t\t\t\tvearth = vi*cosalpha2 + ui*sinalpha2\n",
    "\n",
    "\t\t\tuearth2 = destagger(uearth,0)\n",
    "\t\t\tvearth2 = destagger(vearth,1)\n",
    "\t\t\tvimfc = (uearth2[:,1:]-uearth2[:,0:-1])/dx + (vearth2[1:,:]-vearth2[0:-1,:])/dx\n",
    "        \n",
    "\t\tif first_pass1 == 0:\n",
    "\t\t\tlon2d = s1.variables['XLONG'][-1,:,:]\n",
    "\t\t\tlat2d = s1.variables['XLAT'][-1,:,:]\n",
    "\t\t\tgrid = pyresample.geometry.GridDefinition(lats=lat2d, lons=lon2d)\n",
    "\n",
    "            # Define points\n",
    "\t\t\tswath = pyresample.geometry.SwathDefinition(lons=lon_arr, lats=lat_arr)\n",
    "            # Determine nearest (w.r.t. great circle distance) neighbour in the grid.\n",
    "\t\t\t_, _, index_array, distance_array = pyresample.kd_tree.get_neighbour_info(\n",
    "\t\t\tsource_geo_def=grid, target_geo_def=swath, radius_of_influence=5000,neighbours=1)\n",
    "            # get_neighbour_info() returns indices in the flattened lat/lon grid. Compute the 2D grid indices:\n",
    "\t\t\tcalipso_j1, calipso_i1 = np.unravel_index(index_array, grid.shape)\n",
    "\t\t\tprint (calipso_j1, calipso_i1)\n",
    "\n",
    "\t\t\tfirst_pass1 = 1\n",
    "\n",
    "\t\tcalipso_j = calipso_j1\n",
    "\t\tcalipso_i = calipso_i1\n",
    "        \n",
    "\t\tvimfc2 = vimfc[calipso_j[0]:calipso_j[2],calipso_i[1]:calipso_i[3]]\n",
    "        \n",
    "\t\tvimfc3 = copy.deepcopy(vimfc2)\n",
    "\t\tvimfc3 = vimfc3[0:np.shape(vimfc2)[1],:]\n",
    "\n",
    "\t\tdx_max = np.shape(vimfc3)[1]-1\n",
    "\t\tdy_max = np.shape(vimfc3)[0]-1\n",
    "        \n",
    "\t\tcorr_func_dx = np.zeros([dx_max,np.shape(vimfc3)[0],np.shape(vimfc3)[1]])\n",
    "\t\tfor yy in np.arange(np.shape(vimfc3)[0]): # y dir\n",
    "\t\t\t#print ('Doing yy = ' + str(yy))\n",
    "\t\t\tfor xx in np.arange(np.shape(vimfc3)[1]): # x dir\n",
    "        \n",
    "\t\t\t\tfor xxx in np.arange(1,dx_max+1): # x dir\n",
    "\t\t\t\t\tif xx+xxx < dx_max+1:\n",
    "\t\t\t\t\t\tcorr_func_dx[xxx-1,yy,xx] = vimfc3[yy,xx]*vimfc3[yy,xx+xxx]*xxx\n",
    "\t\t\t\t\telse:\n",
    "\t\t\t\t\t\tcorr_func_dx[xxx-1,yy,xx] = np.nan\n",
    "                        \n",
    "\t\tcorr_func_dy = np.zeros([dy_max,np.shape(vimfc3)[0],np.shape(vimfc3)[1]])\n",
    "\t\tfor xx in np.arange(np.shape(vimfc3)[1]): # x dir\n",
    "\t\t\t#print ('Doing xx = ' + str(xx))\n",
    "\t\t\tfor yy in np.arange(np.shape(vimfc3)[0]): # y dir\n",
    "\n",
    "\t\t\t\tfor yyy in np.arange(1,dy_max+1): # y dir\n",
    "\t\t\t\t\tif yy+yyy < dy_max+1:\n",
    "\t\t\t\t\t\tcorr_func_dy[yyy-1,yy,xx] = vimfc3[yy,xx]*vimfc3[yy+yyy,xx]*yyy\n",
    "\t\t\t\t\telse:\n",
    "\t\t\t\t\t\tcorr_func_dy[yyy-1,yy,xx] = np.nan\n",
    "                        \n",
    "\t\tcorr_func_dx_avg = np.zeros(np.shape(corr_func_dx)[0])\n",
    "\t\tfor ii in np.arange(np.shape(corr_func_dx)[0]):\n",
    "\t\t\tcorr_func_dx_rav = corr_func_dx[ii,:,:].ravel()\n",
    "\t\t\tcorr_func_dx_avg[ii] = np.nanmean(corr_func_dx_rav)\n",
    "    \n",
    "\t\tcorr_func_dy_avg = np.zeros(np.shape(corr_func_dy)[0])\n",
    "\t\tfor ii in np.arange(np.shape(corr_func_dy)[0]):\n",
    "\t\t\tcorr_func_dy_rav = corr_func_dy[ii,:,:].ravel()\n",
    "\t\t\tcorr_func_dy_avg[ii] = np.nanmean(corr_func_dy_rav)\n",
    "\n",
    "\t\tcorr_func_avg_all = np.zeros(np.shape(corr_func_dy)[0])\n",
    "\t\tfor ii in np.arange(np.shape(corr_func_dy)[0]):\n",
    "\t\t\tcorr_func_avg_all[ii] = (corr_func_dx_avg[ii]+corr_func_dy_avg[ii])/2.\n",
    "\n",
    "\t\tif j == 0:\n",
    "\t\t\tcorr_func_avg_all_arr = np.zeros([len(cases),len(corr_func_avg_all)])\n",
    "\t\tcorr_func_avg_all_arr[j,:] = corr_func_avg_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9,9))\n",
    "\n",
    "SMALL_SIZE = 16\n",
    "MEDIUM_SIZE = 20\n",
    "BIGGER_SIZE = 24\n",
    "\n",
    "plt.rc('font', size=SMALL_SIZE)          # controls default text sizes\n",
    "plt.rc('axes', titlesize=MEDIUM_SIZE)     # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels\n",
    "plt.rc('xtick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels\n",
    "plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize\n",
    "plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title\n",
    "\n",
    "xdx = np.arange(1,dx_max+1)\n",
    "cc = ['lightgreen','lightblue','darkgreen','darkblue']\n",
    "labs = [\"Level 2.5\", \"Level 3\", \"Level 2.5 EDMF\",\"YSU\"]\n",
    "\n",
    "for j in np.arange(len(cases)):\n",
    "    minvals = argrelextrema(corr_func_avg_all_arr[j], np.less)\n",
    "    print (minvals[0])\n",
    "    dx2 = 2.*(minvals[0][0]+1)\n",
    "    \n",
    "    #plt.plot(xdx[0:int(dx2)+1],corr_func_avg_all[0:int(dx2)+1],c=cc[j])\n",
    "    plt.plot(2*xdx[0:24],corr_func_avg_all_arr[j,0:24]/np.max(corr_func_avg_all_arr[j,0:24]),c=cc[j],lw=4.)\n",
    "    plt.plot([2*(minvals[0][0]+1),2*(minvals[0][0]+1)],[-1.2,1.02],c=cc[j],ls='--',lw=3.,label=labs[j])\n",
    "    \n",
    "leg = plt.legend(loc='center left', bbox_to_anchor=(0.45, 0.875))\n",
    "leg.get_frame().set_alpha(1.0)\n",
    "plt.xlim(0,50)\n",
    "plt.grid()\n",
    "\n",
    "plt.xlabel('Cell Diameter [km]')\n",
    "plt.ylabel('Normalized Correlation Coefficient [-]')\n",
    "plt.title('2020-03-13 12 UTC')\n",
    "#plt.show()\n",
    "plt.savefig(savedir + \"compare_mar\" + day + \"_mynn_edmf_cell_size\" + \"_\" + active_dom + \"_031320_1200.png\",dpi=300,bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
